# coding:utf-8
import json
from torch.utils.data import DataLoader, Dataset
from ..config.config import Config
from .data_process import TransferData

# 自定义Dataset
class MyDataset(Dataset):
    def __init__(self, data_path):
        super(MyDataset, self).__init__()
        self.config = Config()
        self.dataset = [json.loads(line) for line in open(data_path, encoding='utf8')]

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        content = self.dataset[index]
        text = content['text']
        spo_list = content['spo_list']
        return text, spo_list

class MyDataLoader:
    def __init__(self):
        self.config = Config()
        self.transferData = TransferData()

    def get_data(self):
        # 实例化训练数据集Dataset对象
        train_data = MyDataset(self.config.bert_train_path)

        # # 实例化验证数据集Dataset对象
        dev_data = MyDataset(self.config.bert_dev_path)
        #
        # # 实例化测试数据集Dataset对象
        test_data = MyDataset(self.config.bert_test_path)

        # 实例化训练数据集Dataloader对象
        train_dataloader = DataLoader(dataset=train_data,
                                      batch_size=self.config.bert_batch_size,
                                      shuffle=True,
                                      collate_fn=self.transferData.collate_fn,
                                      drop_last=True)
        # 实例化验证数据集Dataloader对象
        dev_dataloader = DataLoader(dataset=dev_data,
                                    batch_size=self.config.bert_batch_size,
                                    shuffle=True,
                                    collate_fn=self.transferData.collate_fn,
                                    drop_last=True)
        # 实例化测试数据集Dataloader对象
        test_dataloader = DataLoader(dataset=test_data,
                                     batch_size=self.config.bert_batch_size,
                                     shuffle=True,
                                     collate_fn=self.transferData.collate_fn,
                                     drop_last=True)
        return train_dataloader, dev_dataloader, test_dataloader